import pickle

from ensembel.random_forest_train import get_predict

def load_data(file_name):
    '''
    导入待分类的数据集
    :param file_name: 待分类数据文件位置
    :return: test_data
    '''
    f = open(file_name)
    test_data = []
    for line in f.readlines():
        lines = line.strip().split("\t")
        tmp = []
        for x in lines:
            tmp.append(float(x))
        tmp.append(0)
        test_data.append(tmp)
    f.close()
    return test_data

def load_model(result_file, feature_file):
    '''
    导入随机森林模型和每一个分类树中选择的特征
    :param result_file: 随机森林模型文件
    :param feature_file: 分类树特征存储文件
    :return: trees_result: 随机森林模型, trees_fiture: 每一颗分类树选择的特征
    '''

    trees_fiture =[]
    f_fea = open(feature_file)
    for line in f_fea.readlines():
        lines = line.strip().split("\t")
        tmp = []
        for x in lines:
            tmp.append(int(x))
        trees_fiture.append(tmp)
    f_fea.close()

    with open(result_file, "rb") as f:
        trees_result = pickle.load(f)

    return trees_result, trees_fiture

def save_result(data_test, prediction, result_file):
    '''
    保存最终的预测结果
    :param data_test:待预测的数据
    :param prediction: 预测的结果
    :param result_file: 存储最终预测结果的文件名
    '''
    m = len(prediction)
    n = len(data_test[0])

    f_result = open(result_file, "w")
    for i in range(m):
        tmp = []
        for j in range(n - 1):
            tmp.append(str(data_test[i][j]))
        tmp.append(str(prediction[i]))
        f_result.writelines("\t".join(tmp) + "\n")
    f_result.close()

if __name__ == "__main__":

    pwd = "/home/xiefeihong/PycharmProjects/SimpleMachineLearning/static/ensembel/"

    print("------ 1. load test data")
    data_test = load_data(pwd + "test_data.txt")

    print("------ 2. load random forest model")
    trees_result, trees_feature = load_model(pwd + "result_file", pwd + "feature_file")

    print("------ 3. get prediction")
    prediction = get_predict(trees_result, trees_feature, data_test)

    print("------ 4. save result")
    save_result(data_test, prediction, pwd + "final_result")